from asyncio.log import logger
import json
import time
from typing import Literal
import uuid

import falcon
import jwt
from jwcrypto import jwk, jwe
from RedisManager import RedisManager
from SessionTypeDefine import (
    JWT_ENCRPYTION_KEY,
    JWT_SECRET,
    JWT_VERSION,
    SESSION_COOKIE_NAME,
    SESSION_EMPTY_TTL,
    SESSION_FILLED_TTL,
    SSID_COOKIE_NAME,
    JWTPayload,
    JWTPayloadUserPart,
    RedisPayload,
)
from myEncryption import chacha20_encrypt_b64


context = falcon.Context()


class SessionManager:
    def __init__(self) -> None:
        self.db = RedisManager()
        # function alias
        self.set = self.set_one_key
        self.get = self.get_one_key

    async def create(self, res: falcon.Response | None = None) -> tuple[str, dict]:
        ssid = str(uuid.uuid4())
        data = {"createTime": int(time.time())}
        await self.db.set(ssid, json.dumps(data), SESSION_EMPTY_TTL)
        if res:
            self._setSessionCookie_(res, SSID_COOKIE_NAME, ssid, SESSION_EMPTY_TTL)
            self._unsetSessionCookie_(res, SESSION_COOKIE_NAME)
        return ssid, data

    async def fill(
        self,
        req: falcon.Request,
        res: falcon.Response,
        p_data: JWTPayloadUserPart | JWTPayload,
        aes_key: str = ''
    ):
        # 验证会话存在
        if not hasattr(req.context, "ssid") or not req.context.ssid:
            raise ValueError("No valid session found")
        if "exp" not in list(p_data.keys()):
            data: JWTPayload = {
                **p_data,
                "ssid": req.context.ssid,  # 从请求上下文中获取ssid
                "jwtVersion": JWT_VERSION,  # 默认版本
                "exp": int(time.time()) + SESSION_FILLED_TTL,  # 计算过期时间
                "iat": int(time.time()),  # 签发时间
            }
        else:
            data: JWTPayload = p_data
            # if "ssid" not in data:
            #    data["ssid"] = req.context.ssid
            # 这两行好像没用，jwt payload是带有ssid字段的，那基本就会有ssid了

        j: str = jwt.encode(dict(data), JWT_SECRET, algorithm="HS256")
        d: str = chacha20_encrypt_b64(j, JWT_ENCRPYTION_KEY)
        r: RedisPayload = {
            "uid": data["uid"],
            "status": "online",
            "lastActivity": int(time.time()),
            "ipAddress": str(req.remote_addr),
            "userAgent": "",
            "loginMethod": "password",
            "loginTime": int(time.time()),
            "role": data["role"],
            "perms": data["perms"],
            "sessionAesKey": aes_key,
            "createdAt": data["iat"],
            "sessionType": "normal",
        }
        await self.setExpire(data["ssid"], SESSION_FILLED_TTL, "at")
        await self.set_session_data(ssid=data["ssid"], data=r)
        self._setSessionCookie_(res, SSID_COOKIE_NAME, data["ssid"], SESSION_FILLED_TTL)
        self._setSessionCookie_(res, SESSION_COOKIE_NAME, d, SESSION_FILLED_TTL)

    async def delete(self, ssid: str, res: falcon.Response | None = None) -> None:
        _: None = await self.db.delete(ssid)
        if res:
            self._unsetSessionCookie_(res, SSID_COOKIE_NAME)
            self._unsetSessionCookie_(res, SESSION_COOKIE_NAME)

    async def setExpire(
        self, ssid: str, expire: int, operation: Literal["at", "add"] = "at"
    ) -> None:
        try:
            if operation == "add":
                current_ttl = await self.db.getExpire(ssid)
                expire += current_ttl if current_ttl > 0 else 0
            await self.db.setExpire(ssid, expire)
        except Exception as e:
            logger.error(f"Failed to set expire for session {ssid}: {e}")

    async def set_one_key(self, ssid: str, key: str, value: str) -> None:
        try:
            session_data = await self.get_session_data(ssid)
            session_data[key] = value
            await self.db.set(ssid, json.dumps(session_data), "keepttl")
        except Exception as e:
            logger.error(f"Failed to set key {key} for session {ssid}: {e}")

    async def get_one_key(self, ssid: str, key: str) -> str | int | float | bool | None:
        try:
            session_data: dict = await self.get_session_data(ssid)
            return session_data.get(key, "")
        except Exception as e:
            logger.error(f"Failed to get key {key} from session {ssid}: {e}")

    async def get_session_data(self, ssid: str) -> dict:
        try:
            session_data: str = await self.db.get(ssid)
            return json.loads(session_data) if session_data else {}
        except json.JSONDecodeError as e:
            logger.error(f"Invalid session data format for {ssid}: {e}")
            return {}
        except Exception as e:
            logger.error(f"Failed to get session data for {ssid}: {e}")

    async def set_session_data(self, ssid: str, data: str | dict) -> None:
        if isinstance(data, dict):
            d: dict = await self.get_session_data(ssid)
            for k in data:
                d[k] = data[k]
            await self.db.set(ssid, json.dumps(d), "keepttl")
        else:
            await self.db.set(ssid, data, "keepttl")

    async def exists(self, ssid: str) -> bool:
        return await self.db.exists(ssid)

    #########################################################
    def _setSessionCookie_(self, res: falcon.Response, k: str, v: str, max_age=0):
        res.set_cookie(
            name=k,
            value=v,
            expires=None,
            max_age=max_age,
            domain=None,
            path=None,
            secure=None,  # TODO [mark]改成根据配置文件决定是否设置secure
            http_only=True,
            same_site="Strict",
            partitioned=False,
        )

    def _unsetSessionCookie_(
        self,
        res: falcon.Response,
        k: str,
    ):
        res.unset_cookie(
            name=k,
            domain=None,
            path=None,
            same_site="Strict",
        )
